package dnsRequest

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"strings"
)

type Header struct {
	ID              uint64
	Flags           uint64
	QuestionCount   uint64
	AnswerCount     uint64 // RRs is Resource Records
	AuthorityCount  uint64
	AdditionalCount uint64
}

type Query struct {
	QuestionType  uint16
	QuestionClass uint16
}

func (header *Header) SetFlag(qr, opcode, aa, tc, rd, ra, rcode uint16) {
	header.Flags = qr<<15 + opcode<<11 + aa<<10 + tc<<9 + rd<<8 + ra<<7 + rcode
}

func ParseDomainName(domain string) []byte {
	//要将域名解析成相应的格式
	//通过.切分域名，并且将每段的字节数以及内容保存到切片
	//具体结构是长度+内容，长度+内容……最后以0x00结尾
	var (
		buffer   bytes.Buffer
		segments []string = strings.Split(domain, ".")
	)
	for _, seg := range segments {
		binary.Write(&buffer, binary.BigEndian, byte(len(seg)))
		binary.Write(&buffer, binary.BigEndian, []byte(seg))
	}
	binary.Write(&buffer, binary.BigEndian, byte(0x00))

	return buffer.Bytes()
}

// dns 服务器可以查看本机电脑网络配置，获取网络配置中dns服务器即可
// dns 请求端口默认为 53端口
func DigDomain(dnsServerAddr string, domain string) (querys, answers []string) {
	header := Header{}
	// 填充dns首部
	header.ID = 0xFFFF
	header.SetFlag(0, 0, 0, 0, 1, 0, 0)
	header.QuestionCount = 1
	header.AnswerCount = 0
	header.AuthorityCount = 0
	header.AdditionalCount = 0

	query := Query{}
	query.QuestionType = 1 // host address
	query.QuestionClass = 1

	var (
		conn   net.Conn
		err    error
		buffer bytes.Buffer
	)

	//buffer中是我们要发送的数据，里面的内容是DNS首部+查询内容,查询内容包括三个部分、name、type、class
	binary.Write(&buffer, binary.BigEndian, header)
	binary.Write(&buffer, binary.BigEndian, ParseDomainName(domain))
	binary.Write(&buffer, binary.BigEndian, query)

	if conn, err = net.Dial("udp", dnsServerAddr); err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()
	if _, err = conn.Write(buffer.Bytes()); err != nil {
		log.Println(err)
		return
	}
	bytes := make([]byte, 1024)
	n, err := conn.Read(bytes)
	if err != nil {
		log.Println(err)
		return
	}
	return dnsResponseDecode(bytes[:n])
}

func dnsResponseDecode(res []byte) (querys, answers []string) {
	fmt.Println(len(res))
	fmt.Println(res)
	header := res[:12]
	queryNum := uint16(header[4])<<8 + uint16(header[5])
	answerNum := uint16(header[5])<<8 + uint16(header[7])
	data := res[12:]
	index := 0
	querysBytes := make([][]byte, queryNum)
	answersBytes := make([][]byte, answerNum)

	for i := 0; i < int(queryNum); i++ {
		start := index
		l := 0
		for {
			l = int(data[index])
			if l == 0 {
				break
			}
			index += 1 + l
		}
		index += 4
		querysBytes[i] = data[start : index+1]
		index += 1
	}
	if answerNum != 0 {
		for i := 0; i < int(answerNum); i++ {
			start := index
			// name type class 为2字节 time为4字节 datalength 为2 字节
			nums := 2 + 2 + 2 + 4 + 2
			datalenIndex := start + 2 + 2 + 2 + 4
			index = start + nums - 1 + datalenIndex
			answersBytes[i] = data[start : index+1]
			index += 1
		}
	}
	querys = make([]string, queryNum)
	for i, bytes := range querysBytes {
		querys[i] = getQuery(bytes)
	}

	answers = make([]string, answerNum)
	for i, bytes := range answersBytes {
		answers[i] = getAnswerString(bytes)
	}
	return
}

func getQuery(queryBytes []byte) string {
	return getDomain(queryBytes)
}

func getAnswerString(answerBytes []byte) string {
	typ := uint16(answerBytes[2])<<8 + uint16(answerBytes[3])
	datalenIndex := 2 + 2 + 2 + 4
	dataLength := int(uint16(answerBytes[datalenIndex])<<8 + uint16(answerBytes[datalenIndex+1]))
	address := answerBytes[datalenIndex+2 : datalenIndex+2+dataLength]
	if typ == 1 {
		// IP
		return fmt.Sprintf("%d.%d.%d.%d", address[0], address[1], address[2], address[3])
	} else if typ == 5 {
		// CNAME
		return getDomain(address)
	}
	return ""
}

func GetDomain(domainBytes []byte) string {
	return getDomain(domainBytes)
}

func getDomain(domainBytes []byte) string {
	domain := ""
	index := 0
	l := 0
	for {
		if index >= len(domainBytes) {
			break
		}
		l = int(domainBytes[index])
		if l == 0 {
			break
		}
		if index+1+l > len(domainBytes) {
			domain += string(domainBytes[index+1:]) + "."
		} else {
			domain += string(domainBytes[index+1:index+1+l]) + "."
		}
		index += 1 + l
	}
	domain = strings.Trim(domain, ".")
	return domain
}
